package com.hzg.demo.pv;

import com.hzg.util.CommonUtil;
import org.apache.flink.api.common.eventtime.*;
import org.apache.flink.api.java.tuple.Tuple4;
import org.apache.flink.streaming.api.TimeCharacteristic;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.windowing.ProcessAllWindowFunction;
import org.apache.flink.streaming.api.windowing.assigners.SlidingEventTimeWindows;
import org.apache.flink.streaming.api.windowing.time.Time;
import org.apache.flink.streaming.api.windowing.windows.TimeWindow;
import org.apache.flink.streaming.connectors.kafka.FlinkKafkaConsumer;
import org.apache.flink.util.Collector;

import java.util.HashSet;
import java.util.Iterator;
import java.util.Properties;
import java.util.Set;

/**
 * @Author: huangzhigao
 * @Date: 2022/4/19 21:28
 */
public class PvUvExample {
    public static void main(String[] args) {
        StreamExecutionEnvironment env = CommonUtil.getEnv();
        env.setStreamTimeCharacteristic(TimeCharacteristic.EventTime);
        String topic = "test";
        Properties properties = new Properties();
        FlinkKafkaConsumer<UserBehaviorEvent> flinkKafkaConsumer = new FlinkKafkaConsumer<>(topic, new UserBehaviorEventSchema(), properties);
        flinkKafkaConsumer.setStartFromLatest();
        flinkKafkaConsumer.setCommitOffsetsOnCheckpoints(false);
        SerializableTimestampAssigner<UserBehaviorEvent> serializableTimestampAssigner = (SerializableTimestampAssigner<UserBehaviorEvent>) (element, recordTimestamp) -> element.getTs();
//        flinkKafkaConsumer.assignTimestampsAndWatermarks(WatermarkStrategy.<UserBehaviorEvent>forBoundedOutOfOrderness(Duration.ofMinutes(1)).withTimestampAssigner(serializableTimestampAssigner));

        flinkKafkaConsumer.assignTimestampsAndWatermarks(WatermarkStrategy.forGenerator(context -> new PeriodWaterMark()).withTimestampAssigner(context -> new TimeAssigner()));
        DataStreamSource<UserBehaviorEvent> source = env.addSource(flinkKafkaConsumer);
        SingleOutputStreamOperator<Tuple4<Long, Long, Long, Integer>> uvCount = source.windowAll(SlidingEventTimeWindows.of(Time.minutes(10), Time.minutes(5)))
                .allowedLateness(Time.minutes(2))
                .process(new ProcessAllWindowFunction<UserBehaviorEvent, Tuple4<Long, Long, Long, Integer>, TimeWindow>() {

                    @Override
                    public void process(Context context, Iterable<UserBehaviorEvent> elements, Collector<Tuple4<Long, Long, Long, Integer>> out) throws Exception {
                        Long pv = 0L;
                        Set<Integer> userIds = new HashSet();
                        Iterator<UserBehaviorEvent> iterator = elements.iterator();
                        while (iterator.hasNext()) {
                            pv++;
                            userIds.add(iterator.next().getUserId());
                        }
                        TimeWindow window = context.window();
                        out.collect(new Tuple4<>(window.getStart(), window.getEnd(), pv, userIds.size()));
                    }
                });
        uvCount.print();


    }


    /**
     * 无延迟的水位线
     */
    private static class PeriodWaterMark implements WatermarkGenerator<UserBehaviorEvent> {
        private long currentWaterMark = Long.MIN_VALUE;

        @Override
        public void onEvent(UserBehaviorEvent event, long eventTimestamp, WatermarkOutput output) {
            currentWaterMark = event.getTs();
        }

        @Override
        public void onPeriodicEmit(WatermarkOutput output) {
            long effectiveWaterMark = currentWaterMark == Long.MIN_VALUE ? Long.MIN_VALUE : currentWaterMark - 1;
            output.emitWatermark(new Watermark(effectiveWaterMark));
        }
    }


    /**
     * 抽取事件时间
     */
    private static class TimeAssigner implements TimestampAssigner<UserBehaviorEvent> {
        @Override
        public long extractTimestamp(UserBehaviorEvent element, long recordTimestamp) {
            return element.getTs();
        }
    }
}
